import json
import pathlib
import time
from typing import Union

import cv2
import torch
from PIL import Image
from PyQt5.QtCore import QThread, pyqtSignal, Qt, QSize
from PyQt5.QtGui import QMovie, QPixmap, QIcon, QImage
from PyQt5.QtWidgets import QMainWindow, QFileDialog, QMessageBox, QDesktopWidget
from numpy import ndarray
from torch import Tensor

from designer.main_widget import Ui_MainWindow
from designer.threads import VideoThread, ImageThread, CammeraThread, WarmupThread
from model.detect import SingleDetect


class Window(QMainWindow):
    # 信号
    output_print = pyqtSignal(str)
    status_print = pyqtSignal(str)
    show_detect_result = pyqtSignal(object)
    start_loading_signal = pyqtSignal(str)
    stop_loading_signal = pyqtSignal(str)

    def init_slots(self):
        # 输出框与状态栏
        self.output_print.connect(self.print_slot)
        self.status_print.connect(self.status_slot)
        self.show_detect_result.connect(self.show_detect_result_slot)

        # 加载动画
        self.start_loading_signal.connect(self.start_loading_slot)
        self.stop_loading_signal.connect(self.stop_loading_slot)

        # 下拉列表
        self.ui.model_select.activated.connect(self.model_select_activate)

        # 3个按钮
        self.ui.video_button.clicked.connect(self.video_button_clicked)
        self.ui.image_button.clicked.connect(self.image_button_clicked)
        self.ui.cammera_button.clicked.connect(self.cammera_button_clicked)

        # 标签按钮
        self.ui.islabel.clicked.connect(self.radio_button)

        # 数字按钮，重置按钮与滑动条
        self.ui.iou_spinbox.valueChanged.connect(self.iou_spinbox_value_changed)
        self.ui.iou_slider.sliderReleased.connect(self.iou_slider_value_released)
        self.ui.iou_slider.valueChanged.connect(self.iou_slider_value_changed)
        self.ui.conf_spinbox.valueChanged.connect(self.conf_spinbox_value_changed)
        self.ui.conf_slider.valueChanged.connect(self.conf_slider_value_changed)
        self.ui.conf_slider.sliderReleased.connect(self.conf_slider_value_released)
        self.ui.reset_iou_button.clicked.connect(self.reset_iou_button_clicked)
        self.ui.reset_conf_button.clicked.connect(self.reset_conf_button_clicked)


    def init_config(self):
        with open('config.json', encoding="utf8", mode="r") as f:
            self.config = json.load(f)

    def init_ui(self):
        self.init_title_icon()
        # self.init_max_size()
        self.init_model_select()
        self.init_iou_spinbox()
        self.init_conf_spinbox()
        self.init_iou_slider()
        self.init_conf_slider()

    def __init__(self):
        super().__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        # 配置
        self.init_config()

        # 模型
        self.detector: Union[SingleDetect, None] = None
        self.dtype = torch.float32
        self.device = torch.device(self.config.get('device', 'cpu'))

        # 权重
        self.weights_list = []
        self.current_weight_index = 0

        # 模型参数
        self.default_iou = self.config.get('iou', 0.29)
        self.default_conf = self.config.get('conf', 0.29)
        self.iou = self.default_iou
        self.conf = self.default_conf
        self.imgsz = self.config.get('img_size', 640)

        # 线程
        self.warmup_thread = None
        self.video_thread = None
        self.image_thread = None
        self.cammera_thread = None

        self.label_button = True

        # 初始化UI
        self.init_ui()

        # 初始化槽函数
        self.init_slots()

        if len(self.weights_list) > 0:
            self.set_model(self.weights_list[self.current_weight_index])

        self.show()
        self.status_print.emit('一切准备就绪')

    # ======================================================================= #
    # 槽函数

    def start_loading_slot(self, msg):
        if self.ui.show_detect.movie() is None:
            self.set_loading_gif()
        self.ui.show_detect.movie().start()
        # self.status_print.emit(msg)
        self.output_print.emit(msg)
        self.disable_all_button()

    def stop_loading_slot(self, msg):
        if self.ui.show_detect.movie() is not None:
            self.ui.show_detect.movie().stop()
        # self.status_print.emit(msg)
        self.output_print.emit(msg)
        self.enable_all_button()

    def print_slot(self, msg):
        self.ui.output.append(msg)

    def status_slot(self, msg):
        self.ui.status_bar.showMessage(msg, 5000)

    def show_detect_result_slot(self, obj):
        if isinstance(obj, tuple):
            label_img, img0 = obj
            img = label_img if self.label_button else img0
        else:
            img = obj
        if isinstance(img, ndarray):
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(img)
        elif isinstance(img, Tensor):
            # 翻转第0维度
            img = img[[2,1,0],...]
            img = Image.fromarray(img.numpy().transpose(1,2,0))
            
        elif isinstance(img, str):
            if img == '':
                self.ui.show_detect.setPixmap(QPixmap(''))
                return
            img = Image.open(img)
        else:
            print('未知类型')
            print(type(img))
            return
        img = img.convert('RGBA')

        g = self.ui.show_detect.geometry()
        bw, bh = g.width(), g.height()
        iw, ih = img.size
        # 等比缩放
        tw, th = self.get_better_size((bw, bh), (iw, ih), 0.8)
        img = img.resize((tw, th))
        # 居中显示
        self.ui.show_detect.setAlignment(Qt.AlignCenter)
        qimg = QImage(img.tobytes("raw", "RGBA"), img.size[0], img.size[1], QImage.Format_RGBA8888)
        self.ui.show_detect.setPixmap(QPixmap.fromImage(qimg))

    def radio_button(self):
        if self.ui.islabel.isChecked():
            self.label_button = True
        else:
            self.label_button = False

    def model_select_activate(self, index):
        if self.current_weight_index == index:
            return
        self.current_weight_index = index
        self.set_model(self.weights_list[index])

    def iou_spinbox_value_changed(self, value):
        self.ui.iou_slider.setValue(value * 100)
        self.set_model_iou(self.iou)

    def conf_spinbox_value_changed(self, value):
        self.ui.conf_slider.setValue(value * 100)
        self.set_model_conf(self.conf)

    def iou_slider_value_changed(self, value):
        self.ui.iou_spinbox.setValue(value / 100)
        self.iou = value / 100

    def conf_slider_value_changed(self, value):
        self.ui.conf_spinbox.setValue(value / 100)
        self.conf = value / 100

    def conf_slider_value_released(self):
        self.set_model_conf(self.conf)

    def iou_slider_value_released(self):
        self.set_model_iou(self.iou)

    def reset_iou_button_clicked(self):
        self.ui.iou_spinbox.setValue(self.default_iou)
        self.ui.iou_slider.setValue(self.default_iou * 100)
        self.iou = self.default_iou
        self.set_model_iou(self.iou)

    def reset_conf_button_clicked(self):
        self.ui.conf_spinbox.setValue(self.default_conf)
        self.ui.conf_slider.setValue(self.default_conf * 100)
        self.conf = self.default_conf
        self.set_model_conf(self.conf)

    def check_model_error(self):
        if self.detector is None:
            self.output_print.emit('模型未初始化')
            return True
        return False

    def video_button_clicked(self):
        if self.check_model_error(): return
        if self.is_running(self.video_thread):
            self.close_thread(self.video_thread)
            return

        file_name, _ = QFileDialog.getOpenFileName(self, '选择视频', '', '视频文件(*.mp4 *.avi *.mov)')
        if file_name == '':
            return
        self.close_everything()
        self.start_loading_signal.emit('正在打开视频')
        self.video_thread = VideoThread(file_name, self)
        self.video_thread.start()

    def image_button_clicked(self):
        if self.check_model_error(): return
        if self.is_running(self.image_thread):
            return

        file_name, _ = QFileDialog.getOpenFileName(self, '选择图片', '', '图片文件(*.jpg *.png *.jpeg)')
        if file_name == '':
            return
        self.close_everything()
        self.start_loading_signal.emit('正在检测')
        self.image_thread = ImageThread(file_name, self)
        self.image_thread.start()

    def cammera_button_clicked(self):
        if self.check_model_error(): return
        if self.is_running(self.cammera_thread):
            self.close_thread(self.cammera_thread)
            return

        self.close_everything()
        self.start_loading_signal.emit('正在打开摄像头')
        self.cammera_thread = CammeraThread('0', self)
        self.cammera_thread.start()

    # ======================================================================= #
    # UI

    def init_title_icon(self):
        self.setWindowTitle(self.config.get('system-title', "检测系统"))
        self.setWindowIcon(QIcon('assets/icon.png'))
        self.setFixedSize(self.width(), self.height())

    def init_model_select(self):
        weights_dir = pathlib.Path(__file__).parent.parent / 'weights'
        weights_list = [str(weights_dir / x) for x in weights_dir.glob('*.pt')]
        self.weights_list = weights_list
        self.current_weight_index = 0
        # 清除所有选项
        items = [x.split('/')[-1][:-3] for x in weights_list]
        self.ui.model_select.clear()
        self.ui.model_select.addItems(items)

    def init_iou_spinbox(self):
        self.ui.iou_spinbox.setRange(0, 1)
        self.ui.iou_spinbox.setSingleStep(0.01)
        self.ui.iou_spinbox.setValue(self.iou)

    def init_conf_spinbox(self):
        self.ui.conf_spinbox.setRange(0, 1)
        self.ui.conf_spinbox.setSingleStep(0.01)
        self.ui.conf_spinbox.setValue(self.conf)

    def init_iou_slider(self):
        self.ui.iou_slider.setRange(0, 100)
        self.ui.iou_slider.setSingleStep(1)
        self.ui.iou_slider.setValue(int(self.iou * 100))

    def init_conf_slider(self):
        self.ui.conf_slider.setRange(0, 100)
        self.ui.conf_slider.setSingleStep(1)
        self.ui.conf_slider.setValue(int(self.conf * 100))

    def init_max_size(self):
        screen = QDesktopWidget().screenGeometry()
        self.setMaximumSize(screen.width(), screen.height())

    # ======================================================================= #
    # 其他函数
    def set_loading_gif(self):
        self.ui.show_detect.setAlignment(Qt.AlignCenter)
        qm = QMovie('assets/loading.gif')
        g = self.ui.show_detect.geometry()
        bw, bh = g.width(), g.height()
        s = min(bw, bh) * 0.4
        qm.setScaledSize(QSize(s, s))
        self.ui.show_detect.setMovie(qm)

    def set_model(self, path):
        self.warmup_thread = WarmupThread(self, path,
                                          device=self.device, dtype=self.dtype,
                                          iou_thres=self.iou, conf_thres=self.conf,
                                          imgsz=self.imgsz)
        self.warmup_thread.start()

    def set_model_iou(self, iou):
        self.iou = iou
        self.status_print.emit(f'设置模型 iou: {iou}')
        self.detector.iou_thres = iou

    def set_model_conf(self, conf):
        self.conf = conf
        self.status_print.emit(f'设置模型 conf: {conf}')
        self.detector.conf_thres = conf

    def get_better_size(self, bsize, isize, r):
        bw, bh = bsize
        iw, ih = isize
        tw, th = int(r * bw), int(r * bh)
        min_ratio = min(tw / iw, th / ih)
        tw, th = int(min_ratio * iw), int(min_ratio * ih)
        return tw, th

    def disable_all_button(self):
        self.ui.video_button.setEnabled(False)
        self.ui.image_button.setEnabled(False)
        self.ui.cammera_button.setEnabled(False)
        self.ui.reset_iou_button.setEnabled(False)
        self.ui.reset_conf_button.setEnabled(False)
        self.ui.iou_slider.setEnabled(False)
        self.ui.iou_spinbox.setEnabled(False)
        self.ui.conf_slider.setEnabled(False)
        self.ui.conf_spinbox.setEnabled(False)
        self.ui.model_select.setEnabled(False)

    def enable_all_button(self):
        self.ui.video_button.setEnabled(True)
        self.ui.image_button.setEnabled(True)
        self.ui.cammera_button.setEnabled(True)
        self.ui.reset_iou_button.setEnabled(True)
        self.ui.reset_conf_button.setEnabled(True)
        self.ui.iou_slider.setEnabled(True)
        self.ui.iou_spinbox.setEnabled(True)
        self.ui.conf_slider.setEnabled(True)
        self.ui.conf_spinbox.setEnabled(True)
        self.ui.model_select.setEnabled(True)

    def closeEvent(self, a0) -> None:
        msgbox = QMessageBox(QMessageBox.Question, '离开', '确认退出?')
        msgbox.setWindowIcon(self.windowIcon())
        yes = msgbox.addButton('确认', QMessageBox.YesRole)
        msgbox.addButton('取消', QMessageBox.NoRole)
        msgbox.exec()
        if msgbox.clickedButton() == yes:
            self.close_everything()
            a0.accept()
        else:
            a0.ignore()

    def is_running(self, thread: QThread):
        return thread is not None and thread.isRunning()

    def close_thread(self, thread: QThread):
        if thread is not None:
            thread.stop()

    def close_everything(self):
        self.close_thread(self.image_thread)
        self.close_thread(self.video_thread)
        self.close_thread(self.cammera_thread)
        while (self.is_running(self.video_thread) or
               self.is_running(self.image_thread) or
               self.is_running(self.cammera_thread)
        ):
            time.sleep(0.1)
